# import logging
# import os
# import time
# import torch
# import torch.nn as nn
# from utils.meter import AverageMeter
# from utils.metrics import R1_mAP_eval
# from torch.amp import autocast, GradScaler
# import torch.distributed as dist
# import signal


# import logging
# import os
# import time
# import torch
# import torch.nn as nn
# from utils.meter import AverageMeter
# from utils.metrics import R1_mAP_eval
# from torch.amp import autocast, GradScaler
# import torch.distributed as dist
# from torch.utils.tensorboard import SummaryWriter
# import signal
# import sys

#
#
# def do_train(cfg,
#              model,
#              center_criterion,
#              train_loader,
#              val_loader,
#              optimizer,
#              optimizer_center,
#              scheduler,
#              loss_fn,
#              num_query, local_rank):
#     log_period = cfg.SOLVER.LOG_PERIOD
#     checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
#     eval_period = cfg.SOLVER.EVAL_PERIOD
#     output_dir = cfg.OUTPUT_DIR
#     device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
#     epochs = cfg.SOLVER.MAX_EPOCHS
#
#     logger = logging.getLogger("transreid.train")
#     logger.info('start training')
#     logger.info(f"Device: {device}, CUDA available: {torch.cuda.is_available()}")
#     _LOCAL_PROCESS_GROUP = None
#
#     # 添加中断信号处理
#     signal.signal(signal.SIGINT, signal_handler)
#
#     model.to(device)
#     logger.info("Model moved to device")
#     if torch.cuda.device_count() > 1 and cfg.MODEL.DIST_TRAIN:
#         print(f'Using {torch.cuda.device_count()} GPUs for training')
#         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True)
#         logger.info("Initialized DistributedDataParallel")
#
#     loss_meter = AverageMeter()
#     acc_meter = AverageMeter()
#
#     evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
#     scaler = GradScaler('cuda' if torch.cuda.is_available() else 'cpu')
#     logger.info("GradScaler initialized")
#
#     for epoch in range(1, epochs + 1):
#         start_time = time.time()
#         loss_meter.reset()
#         acc_meter.reset()
#         evaluator.reset()
#         scheduler.step(epoch)
#         model.train()
#         logger.info(f"Epoch {epoch} started")
#
#         for n_iter, (img, vid, target_cam, target_view) in enumerate(train_loader):
#             optimizer.zero_grad()
#             if center_criterion is not None and optimizer_center is not None:
#                 optimizer_center.zero_grad()
#
#             img = img.to(device)
#             target = vid.to(device)
#             target_cam = target_cam.to(device)
#             target_view = target_view.to(device)
#
#             with autocast('cuda' if torch.cuda.is_available() else 'cpu'):
#                 score, feat = model(img, target, cam_label=target_cam, view_label=target_view)
#                 loss = loss_fn(score, feat, target, target_cam)
#
#             scaler.scale(loss).backward()
#             scaler.step(optimizer)
#             scaler.update()
#
#             if center_criterion is not None and optimizer_center is not None and 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
#                 for param in center_criterion.parameters():
#                     param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
#                 scaler.step(optimizer_center)
#                 scaler.update()
#
#             if isinstance(score, list):
#                 acc = (score[0].max(1)[1] == target).float().mean()
#             else:
#                 acc = (score.max(1)[1] == target).float().mean()
#
#             loss_meter.update(loss.item(), img.shape[0])
#             acc_meter.update(acc, 1)
#
#             if torch.cuda.is_available():
#                 torch.cuda.synchronize()
#             if (n_iter + 1) % log_period == 0:
#                 logger.info(f"Epoch[{epoch}] Iteration[{n_iter+1}/{len(train_loader)}] Loss: {loss_meter.avg:.3f}, Acc: {acc_meter.avg:.3f}, Base Lr: {scheduler._get_lr(epoch)[0]:.2e}")
#
#         end_time = time.time()
#         time_per_batch = (end_time - start_time) / (n_iter + 1)
#         logger.info(f"Epoch {epoch} done. Time per batch: {time_per_batch:.3f}[s] Speed: {train_loader.batch_size / time_per_batch:.1f}[samples/s]")
#
#         if epoch % checkpoint_period == 0:
#             if cfg.MODEL.DIST_TRAIN:
#                 if dist.get_rank() == 0:
#                     torch.save(model.state_dict(),
#                                os.path.join(output_dir, f"{cfg.MODEL.NAME}_{epoch}.pth"))
#             else:
#                 torch.save(model.state_dict(),
#                            os.path.join(output_dir, f"{cfg.MODEL.NAME}_{epoch}.pth"))
#
#         if epoch % eval_period == 0:
#             model.eval()
#             for n_iter, (img, vid, camid, camids, target_view, _) in enumerate(val_loader):
#                 with torch.no_grad():
#                     img = img.to(device)
#                     camids = camids.to(device)
#                     target_view = target_view.to(device)
#                     feat = model(img, cam_label=camids, view_label=target_view)
#                     evaluator.update((feat, vid, camid))
#             cmc, mAP, _, _, _, _, _ = evaluator.compute()
#             logger.info(f"Validation Results - Epoch: {epoch}")
#             logger.info(f"mAP: {mAP:.1%}")
#             for r in [1, 5, 10]:
#                 logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.1%}")
#             if torch.cuda.is_available():
#                 torch.cuda.empty_cache()
#
#     logger.info("Training finished")
#
# def do_inference(cfg,
#                 model,
#                 val_loader,
#                 num_query):
#     device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
#     logger = logging.getLogger("transreid.test")
#     logger.info("Enter inferencing")
#
#     evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
#     evaluator.reset()
#
#     model.to(device)
#     if torch.cuda.device_count() > 1:
#         print(f'Using {torch.cuda.device_count()} GPUs for inference')
#         model = nn.DataParallel(model)
#
#     model.eval()
#     img_path_list = []
#
#     for n_iter, (img, pid, camid, camids, target_view, imgpath) in enumerate(val_loader):
#         with torch.no_grad():
#             img = img.to(device)
#             camids = camids.to(device)
#             target_view = target_view.to(device)
#             feat = model(img, cam_label=camids, view_label=target_view)
#             evaluator.update((feat, pid, camid))
#             img_path_list.extend(imgpath)
#
#     cmc, mAP, _, _, _, _, _ = evaluator.compute()
#     logger.info("Validation Results ")
#     logger.info(f"mAP: {mAP:.1%}")
#     for r in [1, 5, 10]:
#         logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.1%}")
#     return cmc[0], cmc[4]

import logging
import os
import time
import torch
import torch.nn as nn
from utils.meter import AverageMeter
from utils.metrics import R1_mAP_eval
from torch.amp import autocast, GradScaler
import torch.distributed as dist
from torch.utils.tensorboard import SummaryWriter
import signal
import sys


# 定义 signal_handler 函数以处理 Ctrl+C 中断
def signal_handler(sig, frame):
    logger = logging.getLogger("transreid.train")
    logger.info("Received interrupt signal, saving model and exiting...")
    if not cfg.MODEL.DIST_TRAIN or dist.get_rank() == 0:
        torch.save(model.state_dict(),
                   os.path.join(output_dir, f"{cfg.MODEL.NAME}_interrupted.pth"))
    sys.exit(0)

def do_train(cfg,
             model,
             center_criterion,
             train_loader,
             val_loader,
             optimizer,
             optimizer_center,
             scheduler,
             loss_fn,
             num_query, local_rank,
             camera_num, view_num   # <--- 添加这两个参数
             ):
    log_period = cfg.SOLVER.LOG_PERIOD
    checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
    eval_period = cfg.SOLVER.EVAL_PERIOD
    output_dir = cfg.OUTPUT_DIR
    device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
    epochs = cfg.SOLVER.MAX_EPOCHS

    logger = logging.getLogger("transreid.train")
    logger.info('start training')
    logger.info(f"Device: {device}, CUDA available: {torch.cuda.is_available()}")
    logger.info(f"Training for {epochs} epochs")
    logger.info(f"Validation period: every {eval_period} epochs")
    logger.info(f"Checkpoint period: every {checkpoint_period} epochs")
    logger.info(f"train_loader length: {len(train_loader)}")
    logger.info(f"val_loader length: {len(val_loader)}")

    # 初始化 TensorBoard
    writer = SummaryWriter(os.path.join(output_dir, 'tensorboard'))

    # 添加中断信号处理
    signal.signal(signal.SIGINT, signal_handler)

    model.to(device)
    logger.info("Model moved to device")
    if torch.cuda.device_count() > 1 and cfg.MODEL.DIST_TRAIN:
        print(f'Using {torch.cuda.device_count()} GPUs for training')
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True)
        logger.info("Initialized DistributedDataParallel")

    loss_meter = AverageMeter()
    acc_meter = AverageMeter()

    evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
    scaler = GradScaler('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info("GradScaler initialized")

    for epoch in range(1, epochs + 1):
        start_time = time.time()
        loss_meter.reset()
        acc_meter.reset()
        evaluator.reset()
        scheduler.step(epoch)
        model.train()
        logger.info(f"Epoch {epoch} started")

        for n_iter, (img, vid, target_cam, target_view) in enumerate(train_loader):
            optimizer.zero_grad()
            if center_criterion is not None and optimizer_center is not None:
                optimizer_center.zero_grad()

            img = img.to(device)
            target = vid.to(device)
            target_cam = target_cam.to(device)
            target_view = target_view.to(device)

#添加一段新的测试代码
            # =================== 添加调试代码在这里 ===================
            if n_iter == 0 and epoch == 1:  # 只在第一个batch打印一次，避免刷屏
                print(f"--- Debug Info ---")
                print(f"Camera labels (min, max): {target_cam.min().item()}, {target_cam.max().item()}")
                print(f"View labels (min, max):   {target_view.min().item()}, {target_view.max().item()}")
                # 你还需要知道 camera_num 和 view_num 是多少
                # 这个信息在 train_remake.py 创建模型时有
                # 你可以在调用 do_train 时把它传进来，或者直接去 dataloader 那里看
                print(f"Expected max camera index: {camera_num - 1}")
                print(f"Expected max view index:   {view_num - 1}")
                print(f"--- End Debug Info ---")
            # ==========================================================




            with autocast('cuda' if torch.cuda.is_available() else 'cpu'):
                score, feat = model(img, target, cam_label=target_cam, view_label=target_view)
                loss = loss_fn(score, feat, target, target_cam)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            if center_criterion is not None and optimizer_center is not None and 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
                for param in center_criterion.parameters():
                    param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
                scaler.step(optimizer_center)
                scaler.update()

            if isinstance(score, list):
                acc = (score[0].max(1)[1] == target).float().mean()
            else:
                acc = (score.max(1)[1] == target).float().mean()

            loss_meter.update(loss.item(), img.shape[0])
            acc_meter.update(acc, 1)

            if torch.cuda.is_available():
                torch.cuda.synchronize()
            if (n_iter + 1) % log_period == 0:
                logger.info(f"Epoch[{epoch}] Iteration[{n_iter+1}/{len(train_loader)}] Loss: {loss_meter.avg:.3f}, Acc: {acc_meter.avg:.3f}, Base Lr: {scheduler._get_lr(epoch)[0]:.2e}")
                writer.add_scalar('Loss/train', loss_meter.avg, epoch * len(train_loader) + n_iter)
                writer.add_scalar('Accuracy/train', acc_meter.avg, epoch * len(train_loader) + n_iter)

        end_time = time.time()
        time_per_batch = (end_time - start_time) / (n_iter + 1)
        logger.info(f"Epoch {epoch} done. Time per batch: {time_per_batch:.3f}[s] Speed: {train_loader.batch_size / time_per_batch:.1f}[samples/s]")

        if epoch % checkpoint_period == 0:
            if cfg.MODEL.DIST_TRAIN:
                if dist.get_rank() == 0:
                    torch.save(model.state_dict(),
                               os.path.join(output_dir, f"{cfg.MODEL.NAME}_{epoch}.pth"))
            else:
                torch.save(model.state_dict(),
                           os.path.join(output_dir, f"{cfg.MODEL.NAME}_{epoch}.pth"))
            logger.info(f"Saved checkpoint: {os.path.join(output_dir, f'{cfg.MODEL.NAME}_{epoch}.pth')}")

        if epoch % eval_period == 0:
            model.eval()
            logger.info(f"Starting validation for Epoch {epoch}")
            if len(val_loader) == 0:
                logger.warning("Validation loader is empty, skipping validation")
            else:
                for n_iter, (img, vid, camid, camids, target_view, _) in enumerate(val_loader):
                    with torch.no_grad():
                        img = img.to(device)
                        camids = camids.to(device)
                        target_view = target_view.to(device)
                        feat = model(img, cam_label=camids, view_label=target_view)
                        evaluator.update((feat, vid, camid))
                cmc, mAP, _, _, _, _, _ = evaluator.compute()
                logger.info(f"Validation Results - Epoch: {epoch}")
                logger.info(f"mAP: {mAP:.1%}")
                for r in [1, 5, 10]:
                    logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.1%}")
                writer.add_scalar('mAP/val', mAP, epoch)
                writer.add_scalar('Rank-1/val', cmc[0], epoch)
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    writer.close()
    logger.info("Training finished")

def do_inference(cfg,
                model,
                val_loader,
                num_query):
    device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
    logger = logging.getLogger("transreid.test")
    logger.info("Enter inferencing")

    evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
    evaluator.reset()

    model.to(device)
    if torch.cuda.device_count() > 1:
        print(f'Using {torch.cuda.device_count()} GPUs for inference')
        model = nn.DataParallel(model)

    model.eval()
    img_path_list = []

    for n_iter, (img, pid, camid, camids, target_view, imgpath) in enumerate(val_loader):
        with torch.no_grad():
            img = img.to(device)
            camids = camids.to(device)
            target_view = target_view.to(device)
            feat = model(img, cam_label=camids, view_label=target_view)
            evaluator.update((feat, pid, camid))
            img_path_list.extend(imgpath)

    cmc, mAP, _, _, _, _, _ = evaluator.compute()
    logger.info("Validation Results ")
    logger.info(f"mAP: {mAP:.1%}")
    for r in [1, 5, 10]:
        logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.1%}")
    return cmc[0], cmc[4]